import sys
import pickle
import numpy as np
from matplotlib import pyplot as plt
import tensorflow as tf
import os

np.random.seed(777)
tf.random.set_seed(777)
filename = os.path.basename(__file__)

# parameters
is_use_saver = True
ver = 'v1.2'
L1 = 40
L2 = 80
alpha = 0.001
n_epoch = 50
batch_size = 128

# 正负样本数量
n_positive, n_negative = 2000, 2000

# 生成正样本, 小圆环分布
r_p = 5.0 + tf.random.truncated_normal([n_positive, 1], 0.0, 1.0)
theta_p = tf.random.uniform([n_positive, 1], 0.0, 2 * np.pi)
Xp = tf.concat([r_p * tf.cos(theta_p), r_p * tf.sin(theta_p)], axis=1)
Yp = tf.ones_like(r_p)

# 生成负样本, 大圆环分布
r_n = 8.0 + tf.random.truncated_normal([n_negative, 1], 0.0, 1.0)
theta_n = tf.random.uniform([n_negative, 1], 0.0, 2 * np.pi)
Xn = tf.concat([r_n * tf.cos(theta_n), r_n * tf.sin(theta_n)], axis=1)
Yn = tf.zeros_like(r_n)

# 汇总样本
X = tf.concat([Xp, Xn], axis=0)
Y = tf.concat([Yp, Yn], axis=0)

# 可视化
plt.figure(figsize=(6, 6))
plt.scatter(Xp[:, 0].numpy(), Xp[:, 1].numpy(), c="r")
plt.scatter(Xn[:, 0].numpy(), Xn[:, 1].numpy(), c="g")
plt.legend(["positive", "negative"])
plt.show()


# 构建数据管道迭代器
def data_iter(features, labels, batch_size=8):
    num_examples = len(features)
    indices = list(range(num_examples))
    np.random.shuffle(indices)  # 样本的读取顺序是随机的
    for i in range(0, num_examples, batch_size):
        indexs = indices[i: min(i + batch_size, num_examples)]
        yield tf.gather(features, indexs), tf.gather(labels, indexs)


# parameters deduced
m, n = X.shape
n_cls = len(np.unique(Y.numpy()))


class DNNModel(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.w1 = tf.Variable(tf.random.truncated_normal([n, L1]), dtype=tf.float32, name='w1')
        self.b1 = tf.Variable(tf.zeros([1, L1]), dtype=tf.float32, name='b1')
        self.w2 = tf.Variable(tf.random.truncated_normal([L1, L2]), dtype=tf.float32, name='w2')
        self.b2 = tf.Variable(tf.zeros([1, L2]), dtype=tf.float32, name='b2')
        self.w3 = tf.Variable(tf.random.truncated_normal([L2, 1]), dtype=tf.float32, name='w3')
        self.b3 = tf.Variable(tf.zeros([1, 1]), dtype=tf.float32, name='b3')
        self.vars = [
            self.w1, self.b1,
            self.w2, self.b2,
            self.w3, self.b3,
        ]

    def __call__(self, x):
        x1 = tf.nn.relu(x @ self.w1 + self.b1)
        x2 = tf.nn.relu(x1 @ self.w2 + self.b2)
        h = tf.nn.sigmoid(x2 @ self.w3 + self.b3)
        return h

    @staticmethod
    def loss_func(y_true, y_pred):
        eps = 1e-7
        y_pred = tf.clip_by_value(y_pred, eps, 1. - eps)  # ATTENTION
        bce = - y_true * tf.math.log(y_pred) - (1 - y_true) * tf.math.log(1 - y_pred)
        return tf.reduce_mean(bce)

    @staticmethod
    def metric_func(y_true, y_pred):
        acc = tf.reduce_mean(
            tf.cast(
                tf.equal(
                    y_true > 0.5,
                    y_pred > 0.5
                ),
                tf.float32
            )
        )
        return acc


save_path = './_save/' + filename + '_' + ver + '.tmp.dat'
if is_use_saver and os.path.exists(save_path):
    with open(save_path, 'br') as f:
        model = pickle.load(f)
else:
    model = DNNModel('my_logistic_regr')


    @tf.function
    def train_step(model, x, y, alpha=0.01):
        with tf.GradientTape() as tape:
            tape.watch(model.vars)
            pred = model(x)
            loss = model.loss_func(y, pred)
        grads = tape.gradient(loss, model.vars)

        for p, dloss_dp in zip(model.vars, grads):
            p.assign(p - alpha * dloss_dp)

        metric = model.metric_func(y, pred)
        return loss, metric


    def train_model(model, n_epoch):
        writer = tf.summary.create_file_writer('./_log/' + filename)
        g_step = -1
        for epoch in range(n_epoch):
            i = -1
            for bx, by in data_iter(X, Y, batch_size):
                i += 1
                g_step += 1
                loss, metric = train_step(model, bx, by)
                with writer.as_default():
                    tf.summary.scalar('loss', loss, g_step)
                    tf.summary.scalar('metric', metric, g_step)
            print(f'g_step#{g_step + 1}: epoch#{epoch + 1}: loss = {loss}, metric = {metric}')


    train_model(model, n_epoch)
    with open(save_path, 'bw') as f:
        pickle.dump(model, f)

# 结果可视化
fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(12, 5))
ax1.scatter(Xp[:, 0], Xp[:, 1], c="r")
ax1.scatter(Xn[:, 0], Xn[:, 1], c="g")
ax1.legend(["positive", "negative"]);
ax1.set_title("y_true");

Xp_pred = tf.boolean_mask(X, tf.squeeze(model(X) >= 0.5), axis=0)
Xn_pred = tf.boolean_mask(X, tf.squeeze(model(X) < 0.5), axis=0)

ax2.scatter(Xp_pred[:, 0], Xp_pred[:, 1], c="r")
ax2.scatter(Xn_pred[:, 0], Xn_pred[:, 1], c="g")
ax2.legend(["positive", "negative"]);
ax2.set_title("y_pred");
plt.show()
